#include "qcow2.h"
#include "utils.h"
#include "osdep.h"
#include "bswap.h"
#include <string.h>
#include <stdarg.h>
#include <endian.h>
#include <byteswap.h>
#include <stdint.h>
#include <algorithm>
#include <vector>
#include <map>

/* not used
// 文件操作包装（带错误处理）
static int qcow2_file_seek(FILE *f, uint64_t offset) {
    if (fseek(f, (long)offset, SEEK_SET) != 0) {
        report_error("Seek failed to offset 0x%lx", offset);
        return QCOW2_ERROR_IO;
    }
    return 0;
}

static int qcow2_file_read(FILE *f, void *buf, size_t size) {
    if (fread(buf, 1, size, f) != size) {
        report_error("Read failed: requested %zu bytes", size);
        return QCOW2_ERROR_IO;
    }
    return 0;
}

static int qcow2_file_write(FILE *f, const void *buf, size_t size) {
    if (fwrite(buf, 1, size, f) != size) {
        report_error("Write failed: attempted %zu bytes", size);
        return QCOW2_ERROR_IO;
    }
    return 0;
}
*/
/*
// 对齐偏移量到簇大小
static uint64_t align_offset(uint64_t offset, int cluster_size)
{
    return (offset + cluster_size - 1) & ~(cluster_size - 1);
}

// 填充文件到指定偏移量
static int pad_file_to_offset(FILE *f, uint64_t target_offset)
{
    uint64_t current_offset = ftell(f);

    if (current_offset < target_offset)
    {
        uint64_t padding_size = target_offset - current_offset;
        uint8_t *padding = (uint8_t *)qemu_mallocz(padding_size);

        if (fwrite(padding, 1, padding_size, f) != padding_size)
        {
            qemu_free(padding);
            return -1;
        }

        qemu_free(padding);
    }

    return 0;
}
*/
// 探测 QCow2 文件
int qcow2_probe(const uint8_t *buf, int buf_size)
{
    const QCowHeaderV3 *header = (const QCowHeaderV3 *)buf;

    if (buf_size >= (int)sizeof(QCowHeaderV3) &&
        be32_to_cpu(header->magic) == QCOW_MAGIC &&
        (be32_to_cpu(header->version) == QCOW_VERSION ||
         be32_to_cpu(header->version) == QCOW_VERSION3))
    {
        return 1;
    }
    return 0;
}

// 设置 backing file
int qcow2_set_backing_file(BlockDriverState *bs, const char *backing_file)
{
    QCow2State *s = (QCow2State *)bs->opaque;
    if (!s || !backing_file)
        return QCOW2_ERROR_INVALID;

    FILE *f = fopen(bs->filename, "rb+");
    if (!f)
        return QCOW2_ERROR_IO;

    // 读取头文件
    QCowHeaderV3 header;
    if (fread(&header, 1, sizeof(header), f) != sizeof(header))
    {
        fclose(f);
        return QCOW2_ERROR_IO;
    }

    // 查找文件末尾，寻找存放 backing file 的位置
    fseek(f, 0, SEEK_END);
    uint64_t backing_file_offset = ftell(f);

    // 写入 backing file 名称
    size_t backing_file_len = strlen(backing_file);
    if (fwrite(backing_file, 1, backing_file_len + 1, f) != backing_file_len + 1)
    {
        fclose(f);
        return QCOW2_ERROR_IO;
    }

    // 更新头文件中的 backing file 信息
    header.backing_file_offset = cpu_to_be64(backing_file_offset);
    header.backing_file_size = cpu_to_be32(backing_file_len);

    // 写回头文件
    fseek(f, 0, SEEK_SET);
    if (fwrite(&header, 1, sizeof(header), f) != sizeof(header))
    {
        fclose(f);
        return QCOW2_ERROR_IO;
    }

    fclose(f);

    report_info("Set backing file to: %s (offset: 0x%lx)", backing_file, backing_file_offset);
    return 0;
}

// 获取 backing file
char *qcow2_get_backing_file(BlockDriverState *bs)
{
    QCow2State *s = (QCow2State *)bs->opaque;
    if (!s)
        return NULL;

    FILE *f = fopen(bs->filename, "rb");
    if (!f)
        return NULL;

    // 读取头文件
    QCowHeaderV3 header;
    if (fread(&header, 1, sizeof(header), f) != sizeof(header))
    {
        fclose(f);
        return NULL;
    }

    // 字节序转换
    uint64_t backing_file_offset = be64_to_cpu(header.backing_file_offset);
    uint32_t backing_file_size = be32_to_cpu(header.backing_file_size);

    if (backing_file_offset == 0 || backing_file_size == 0)
    {
        fclose(f);
        return NULL;
    }

    // 读取 backing file 名称
    char *backing_file = (char *)qemu_malloc(backing_file_size + 1);
    if (fseek(f, backing_file_offset, SEEK_SET) != 0 ||
        fread(backing_file, 1, backing_file_size, f) != backing_file_size)
    {
        qemu_free(backing_file);
        fclose(f);
        return NULL;
    }

    backing_file[backing_file_size] = '\0';
    fclose(f);

    return backing_file;
}

// 创建 QCow2 文件
int qcow2_create(const char *filename, int64_t size,
                 const char *backing_file, int cluster_size, int version)
{
    if (version != QCOW_VERSION && version != QCOW_VERSION3)
    {
        report_error("Unsupported version: %d (supported: 2, 3)", version);
        return QCOW2_ERROR_INVALID;
    }

    if (size <= 0)
    {
        report_error("Invalid size: %ld", size);
        return QCOW2_ERROR_INVALID;
    }

    if (file_exists(filename))
    {
        report_error("File already exists: %s", filename);
        return QCOW2_ERROR_EXIST;
    }

    if (cluster_size <= 0)
    {
        cluster_size = 65536; // 默认64K
    }

    // 验证簇大小是2的幂
    if (!is_power_of_two((uint64_t)cluster_size))
    {
        report_error("Cluster size must be power of 2: %d", cluster_size);
        return QCOW2_ERROR_INVALID;
    }

    FILE *f;
    int cluster_bits = 0;

    // 计算簇大小对应的位数
    while ((1 << cluster_bits) < cluster_size)
    {
        cluster_bits++;
    }

    f = fopen(filename, "wb");
    if (!f)
    {
        report_error("Failed to create file: %s", filename);
        return QCOW2_ERROR_IO;
    }
    size_t cluster_size_t = (size_t)cluster_size;
    if (version == QCOW_VERSION)
    {
        // Version 2 创建
        QCowHeaderV3 header;
        uint64_t *l1_table;
        uint32_t l1_size;

        memset(&header, 0, sizeof(header));
        header.magic = cpu_to_be32(QCOW_MAGIC);
        header.version = cpu_to_be32(QCOW_VERSION);
        header.cluster_bits = cpu_to_be32((uint32_t)cluster_bits);
        header.size = cpu_to_be64((uint64_t)size);
        header.crypt_method = cpu_to_be32(QCOW_CRYPT_NONE);

        // 计算 L1 表大小
        int l2_bits = cluster_bits - 3;
        l1_size = (uint32_t)((size + (1LL << (cluster_bits + l2_bits)) - 1) >> (cluster_bits + l2_bits));

        // 确保 L1 表至少有一个条目
        if (l1_size == 0)
            l1_size = 1;

        header.l1_size = cpu_to_be32(l1_size);
        header.l1_table_offset = cpu_to_be64(sizeof(header));

        // 写入头部
        if (fwrite(&header, 1, sizeof(header), f) != sizeof(header))
        {
            report_error("Failed to write header");
            fclose(f);
            return QCOW2_ERROR_IO;
        }

        // 写入空的 L1 表
        l1_table = (uint64_t *)qemu_mallocz(l1_size * sizeof(uint64_t));
        // 字节序转换
        for (uint32_t i = 0; i < l1_size; i++)
        {
            l1_table[i] = cpu_to_be64(0); // 全部初始化为0
        }

        if (fwrite(l1_table, sizeof(uint64_t), l1_size, f) != l1_size)
        {
            report_error("Failed to write L1 table");
            qemu_free(l1_table);
            fclose(f);
            return QCOW2_ERROR_IO;
        }
        qemu_free(l1_table);

        // 如果指定了 backing file，在文件末尾写入
        if (backing_file && backing_file[0] != '\0')
        {
            uint64_t backing_offset = (uint64_t)ftell(f);
            size_t backing_len = strlen(backing_file);

            // 写入 backing file 名称
            if (fwrite(backing_file, 1, backing_len + 1, f) != backing_len + 1)
            {
                report_error("Failed to write backing file name");
                fclose(f);
                return QCOW2_ERROR_IO;
            }

            // 更新头文件中的 backing file 信息
            header.backing_file_offset = cpu_to_be64(backing_offset);
            header.backing_file_size = cpu_to_be32((uint32_t)backing_len);

            // 重新写回头文件
            fseek(f, 0, SEEK_SET);
            if (fwrite(&header, 1, sizeof(header), f) != sizeof(header))
            {
                report_error("Failed to update header with backing file info");
                fclose(f);
                return QCOW2_ERROR_IO;
            }
        }

        fclose(f);
    }
    else
    {

        // char cmd[512];
        // sprintf(cmd, "qemu-img create -f qcow2 -b %s %s %lld", backing_file, filename, size);
        // system(cmd);

        QCowHeaderV3 header;
        memset(&header, 0, sizeof(header));

        header.magic = cpu_to_be32(QCOW_MAGIC);
        header.version = cpu_to_be32(QCOW_VERSION3);

        // Backing file 信息
        uint64_t backing_file_offset = 0;
        uint32_t backing_file_size = 0;

        if (backing_file && backing_file[0] != '\0')
        {
            backing_file_size = (uint32_t)strlen(backing_file);
            // QEMU 将 backing file 放在扩展数据区域
            backing_file_offset = 0x138; // 这是 QEMU 实际使用的位置
        }

        header.backing_file_offset = cpu_to_be64(backing_file_offset);
        header.backing_file_size = cpu_to_be32(backing_file_size);

        header.cluster_bits = cpu_to_be32((uint32_t)cluster_bits);
        header.size = cpu_to_be64((uint64_t)size);
        header.crypt_method = cpu_to_be32(0);

        // 计算 L1 表大小
        int l2_bits = cluster_bits - 3;
        uint64_t l1_entries = (uint64_t)((size + (1LL << (cluster_bits + l2_bits)) - 1) >> (cluster_bits + l2_bits));
        if (l1_entries == 0) l1_entries = 1;
        if (l1_entries > (1U << 31)) {
            report_error("L1 table too large");
            fclose(f);
            return QCOW2_ERROR_INVALID;
        }
        uint32_t l1_size = (uint32_t)l1_entries;

        header.l1_size = cpu_to_be32(l1_size);
        header.l1_table_offset = cpu_to_be64(0x10000); // 簇 1

        header.refcount_table_offset = cpu_to_be64(0x20000); // 簇 2
        header.refcount_table_clusters = cpu_to_be32(1);

        header.nb_snapshots = cpu_to_be32(0);
        header.snapshots_offset = cpu_to_be64(0);

        header.incompatible_features = cpu_to_be64(0);
        header.compatible_features = cpu_to_be64(0);
        header.autoclear_features = cpu_to_be64(0);

        header.refcount_order = cpu_to_be32(4);
        header.header_length = cpu_to_be32(104); // **关键：使用 104 而不是 72**

        report_info("Creating QCow2 v3 with header_length=%u",
                   be32_to_cpu(header.header_length));

        // 1. 写入头部
        if (fwrite(&header, 1, sizeof(header), f) != sizeof(header))
        {
            report_error("Failed to write header");
            fclose(f);
            return QCOW2_ERROR_IO;
        }

        if (backing_file && backing_file[0] != '\0')
        {
            // 填充到 0x138 位置
            long current_pos = sizeof(header);
            if (current_pos < 0x138)
            {
                uint64_t padding_size = (uint64_t)(0x138 - current_pos);
                uint8_t *padding = (uint8_t *)qemu_mallocz(padding_size);
                if (!padding)
                {
                    report_error("Failed to allocate padding");
                    fclose(f);
                    return QCOW2_ERROR_IO;
                }
                if (fwrite(padding, 1, padding_size, f) != padding_size)
                {
                    report_error("Failed to pad to extension data");
                    qemu_free(padding);
                    fclose(f);
                    return QCOW2_ERROR_IO;
                }
                qemu_free(padding);
                report_info("Padded to 0x138 for extension data");
            }

            // 在 0x138 写入 backing file 名称
            if (fwrite(backing_file, 1, backing_file_size, f) != backing_file_size)
            {
                report_error("Failed to write backing file name");
                fclose(f);
                return QCOW2_ERROR_IO;
            }
            report_info("Backing file written at 0x%lx: '%s'",
                       0x138L, backing_file);
        }

        // 填充到第一个簇 (0x10000)
        long current_pos = ftell(f);
        if (current_pos < 0x10000)
        {
            uint64_t padding_size = (uint64_t)(0x10000 - current_pos);
            uint8_t *padding = (uint8_t *)qemu_mallocz(padding_size);
            if (!padding)
            {
                report_error("Failed to allocate padding");
                fclose(f);
                return QCOW2_ERROR_IO;
            }
            if (fwrite(padding, 1, padding_size, f) != padding_size)
            {
                report_error("Failed to pad to cluster 1");
                qemu_free(padding);
                fclose(f);
                return QCOW2_ERROR_IO;
            }
            qemu_free(padding);
        }

        // 写入 L1 表
        uint64_t *l1_table = (uint64_t *)qemu_mallocz(l1_size * sizeof(uint64_t));
        if (!l1_table)
        {
            report_error("Failed to allocate L1 table");
            fclose(f);
            return QCOW2_ERROR_IO;
        }

        for (uint32_t i = 0; i < l1_size; i++)
        {
            l1_table[i] = cpu_to_be64(0);
        }

        if (fwrite(l1_table, 1, l1_size * sizeof(uint64_t), f) != l1_size * sizeof(uint64_t))
        {
            report_error("Failed to write L1 table");
            qemu_free(l1_table);
            fclose(f);
            return QCOW2_ERROR_IO;
        }
        qemu_free(l1_table);

        // 填充 L1 表到簇边界
        current_pos = (long)(0x10000 + l1_size * sizeof(uint64_t));
        if (current_pos < 0x20000)
        {
            uint64_t padding_size = (uint64_t)(0x20000 - current_pos);
            uint8_t *padding = (uint8_t *)qemu_mallocz(padding_size);
            if (!padding)
            {
                report_error("Failed to allocate padding");
                fclose(f);
                return QCOW2_ERROR_IO;
            }
            if (fwrite(padding, 1, padding_size, f) != padding_size)
            {
                report_error("Failed to pad after L1 table");
                qemu_free(padding);
                fclose(f);
                return QCOW2_ERROR_IO;
            }
            qemu_free(padding);
        }

        // 写入引用计数表
        uint64_t *refcount_table = (uint64_t *)qemu_mallocz((size_t)cluster_size);
        if (!refcount_table)
        {
            report_error("Failed to allocate refcount table");
            fclose(f);
            return QCOW2_ERROR_IO;
        }

        refcount_table[0] = cpu_to_be64(0x30000);

        if (fwrite(refcount_table, 1, cluster_size_t, f) != cluster_size_t)
        {
            report_error("Failed to write refcount table");
            qemu_free(refcount_table);
            fclose(f);
            return QCOW2_ERROR_IO;
        }
        qemu_free(refcount_table);

        // 写入引用计数块
        uint16_t *refcount_block = (uint16_t *)qemu_mallocz((size_t)cluster_size);
        if (!refcount_block)
        {
            report_error("Failed to allocate refcount block");
            fclose(f);
            return QCOW2_ERROR_IO;
        }

        refcount_block[0] = cpu_to_be16(1);
        refcount_block[1] = cpu_to_be16(1);
        refcount_block[2] = cpu_to_be16(1);
        refcount_block[3] = cpu_to_be16(1);

        if (fwrite(refcount_block, 1, cluster_size_t, f) != cluster_size_t)
        {
            report_error("Failed to write refcount block");
            qemu_free(refcount_block);
            fclose(f);
            return QCOW2_ERROR_IO;
        }
        qemu_free(refcount_block);

        fclose(f);
    }

    report_info("QCow2 file created: %s, version=%d, size=%ld, clustersize=%d%s",
                filename, version, size, cluster_size,
                (backing_file && backing_file[0] != '\0') ? ", with backing file" : "");

    return 0;
}

// 打开 QCow2 文件
int qcow2_open(BlockDriverState *bs, const char *filename)
{
    QCowHeaderV3 header;
    QCow2State *s;
    FILE *f;
    size_t ret;
    uint32_t version;

    report_info("Opening QCow2 file: %s", filename);

    f = fopen(filename, "rb");
    if (!f)
    {
        report_error("Failed to open file: %s", filename);
        return QCOW2_ERROR_IO;
    }

    // 先读取基础头文件
    ret = fread(&header, 1, sizeof(header), f);
    if (ret != sizeof(header))
    {
        report_error("Failed to read QCow2 header");
        fclose(f);
        return QCOW2_ERROR_IO;
    }

    // 字节序转换基础字段
    header.magic = be32_to_cpu(header.magic);
    header.version = be32_to_cpu(header.version);
    version = header.version;

    if (header.magic != QCOW_MAGIC)
    {
        report_error("Not a QCow2 file: magic=0x%x, expected=0x%x",
                     header.magic, QCOW_MAGIC);
        fclose(f);
        return QCOW2_ERROR_INVALID;
    }

    if (header.version != QCOW_VERSION && header.version != QCOW_VERSION3)
    {
        report_error("Unsupported QCow2 version: %d", header.version);
        fclose(f);
        return QCOW2_ERROR_INVALID;
    }

    // 根据版本读取完整的头文件
    if (version == QCOW_VERSION3)
    {
        // 重新读取 Version 3 头文件
        fseek(f, 0, SEEK_SET);
        ret = fread(&header, 1, sizeof(header), f);
        if (ret != sizeof(header))
        {
            report_error("Failed to read QCow2 v3 header");
            fclose(f);
            return QCOW2_ERROR_IO;
        }

        // 字节序转换 Version 3 头文件
        header.magic = be32_to_cpu(header.magic);
        header.version = be32_to_cpu(header.version);
        header.backing_file_offset = be64_to_cpu(header.backing_file_offset);
        header.backing_file_size = be32_to_cpu(header.backing_file_size);
        header.cluster_bits = be32_to_cpu(header.cluster_bits);
        header.size = be64_to_cpu(header.size);
        header.crypt_method = be32_to_cpu(header.crypt_method);
        header.l1_size = be32_to_cpu(header.l1_size);
        header.l1_table_offset = be64_to_cpu(header.l1_table_offset);
        header.refcount_table_offset = be64_to_cpu(header.refcount_table_offset);
        header.refcount_table_clusters = be32_to_cpu(header.refcount_table_clusters);
        header.nb_snapshots = be32_to_cpu(header.nb_snapshots);
        header.snapshots_offset = be64_to_cpu(header.snapshots_offset);
        header.incompatible_features = be64_to_cpu(header.incompatible_features);
        header.compatible_features = be64_to_cpu(header.compatible_features);
        header.autoclear_features = be64_to_cpu(header.autoclear_features);
        header.refcount_order = be32_to_cpu(header.refcount_order);
        header.header_length = be32_to_cpu(header.header_length);
    }
    else
    {
        // Version 2 头文件继续转换
        header.backing_file_offset = be64_to_cpu(header.backing_file_offset);
        header.backing_file_size = be32_to_cpu(header.backing_file_size);
        header.cluster_bits = be32_to_cpu(header.cluster_bits);
        header.size = be64_to_cpu(header.size);
        header.crypt_method = be32_to_cpu(header.crypt_method);
        header.l1_size = be32_to_cpu(header.l1_size);
        header.l1_table_offset = be64_to_cpu(header.l1_table_offset);
        header.refcount_table_offset = be64_to_cpu(header.refcount_table_offset);
        header.refcount_table_clusters = be32_to_cpu(header.refcount_table_clusters);
        header.nb_snapshots = be32_to_cpu(header.nb_snapshots);
        header.snapshots_offset = be64_to_cpu(header.snapshots_offset);
    }

    // 初始化 QCow2 状态
    s = (QCow2State *)qemu_mallocz(sizeof(QCow2State));
    s->version = version;
    s->cluster_bits = header.cluster_bits;
    s->cluster_size = 1 << header.cluster_bits;
    s->cluster_sectors = s->cluster_size / 512;
    s->l2_bits = s->cluster_bits - 3;
    s->l2_size = 1 << s->l2_bits;
    s->l1_size = header.l1_size;
    s->crypt_method = header.crypt_method;
    s->refcount_order = (version == QCOW_VERSION3) ? header.refcount_order : 4;
    s->cluster_offset_mask = ~(QCOW_OFLAG_COMPRESSED | QCOW_OFLAG_COPIED);

    // Version 3 特性
    if (version == QCOW_VERSION3)
    {
        s->incompatible_features = header.incompatible_features;
        s->compatible_features = header.compatible_features;
        s->autoclear_features = header.autoclear_features;
    }

    bs->total_sectors = header.size / 512;
    bs->cluster_size = s->cluster_size;
    bs->opaque = s;
    bs->filename = qemu_strdup(filename);
    s->bs = bs;

    // 读取 L1 表
    s->l1_table_offset = header.l1_table_offset;
    if (qcow2_read_l1_table(s, f) < 0)
    {
        report_error("Failed to read L1 table");
        qcow2_close(bs);
        fclose(f);
        return QCOW2_ERROR_CORRUPT;
    }

    // 初始化缓存
    s->l2_cache_size = 64;
    s->l2_cache = (uint64_t *)qemu_mallocz(s->l2_cache_size * s->l2_size * sizeof(uint64_t));
    s->l2_cache_index = 0;
    for (int i = 0; i < s->l2_cache_size; i++)
    {
        s->l2_cache_valid[i] = false;
    }

    s->cluster_cache = (uint8_t *)qemu_mallocz(s->cluster_size);
    s->cluster_cache_offset = -1;

    fclose(f);

    report_info("QCow2 file opened successfully: version=%d, size=%lu, clustersize=%d",
                version, header.size, s->cluster_size);

    return 0;
}

// L1 表读取函数
int qcow2_read_l1_table(QCow2State *s, FILE *f)
{
    if (s->l1_size == 0)
    {
        report_error("L1 table size is 0");
        return -1;
    }

    s->l1_table = (uint64_t *)qemu_malloc(s->l1_size * sizeof(uint64_t));

    if (fseek(f, s->l1_table_offset, SEEK_SET) != 0)
    {
        report_error("Failed to seek to L1 table at offset 0x%lx", s->l1_table_offset);
        return -1;
    }

    if (fread(s->l1_table, sizeof(uint64_t), s->l1_size, f) != (size_t)s->l1_size)
    {
        report_error("Failed to read L1 table");
        return -1;
    }

    // 字节序转换
    for (int i = 0; i < s->l1_size; i++)
    {
        s->l1_table[i] = be64_to_cpu(s->l1_table[i]);
    }

    report_info("L1 table loaded: %d entries", s->l1_size);
    return 0;
}

int qcow2_read_l2_table(QCow2State *s, FILE *f, uint64_t l2_offset,
                        uint64_t **l2_table)
{
    static int cache_index = 0;

    // 简单缓存实现
    *l2_table = &s->l2_cache[cache_index * s->l2_size];

    if (fseek(f, (long)l2_offset, SEEK_SET) != 0)
    {
        report_error("Failed to seek to L2 table");
        return -1;
    }

    size_t read_size = (size_t)s->l2_size;
    if (fread(*l2_table, sizeof(uint64_t), read_size, f) != read_size)
    {
        report_error("Failed to read L2 table");
        return -1;
    }

    // 字节序转换
    for (int i = 0; i < s->l2_size; i++)
    {
        (*l2_table)[i] = be64_to_cpu((*l2_table)[i]);
    }

    s->l2_cache_offsets[cache_index] = l2_offset;
    cache_index = (cache_index + 1) % 16;

    return 0;
}

// 关闭 QCow2 文件
int qcow2_close(BlockDriverState *bs)
{
    QCow2State *s = (QCow2State *)bs->opaque;

    if (s)
    {
        if (s->l1_table)
        {
            qemu_free(s->l1_table);
        }
        if (s->l2_cache)
        {
            qemu_free(s->l2_cache);
        }
        if (s->cluster_cache)
        {
            qemu_free(s->cluster_cache);
        }
        if (s->refcount_table)
        {
            qemu_free(s->refcount_table);
        }
        // 清理引用计数块缓存
        if (s && !s->refcount_block_cache.empty())
        {
            report_info("Cleaning up refcount block cache, size: %zu", s->refcount_block_cache.size());

            // 使用迭代器而不是范围for循环，更安全
            auto it = s->refcount_block_cache.begin();
            while (it != s->refcount_block_cache.end())
            {
                if (it->second)
                {
                    if (it->second->counts)
                    {
                        qemu_free(it->second->counts);
                        it->second->counts = NULL;
                    }
                    qemu_free(it->second);
                    it->second = NULL;
                }
                auto current = it;
                it++;
                s->refcount_block_cache.erase(current);
            }
            s->refcount_block_cache.clear();
            report_info("Refcount block cache cleared");
        }
        else
        {
            report_info("No refcount block cache to clean up");
        }
        qemu_free(s);
    }
    if (bs->filename)
    {
        qemu_free(bs->filename);
        bs->filename = NULL;
    }
    report_info("QCow2 file closed");
    return 0;
}

// 获取文件长度
int64_t qcow2_getlength(BlockDriverState *bs)
{
    return bs->total_sectors * 512;
}

// 读取数据
int qcow2_read(BlockDriverState *bs, int64_t offset, uint8_t *buf, int64_t bytes)
{
    QCow2State *s = (QCow2State *)bs->opaque;
    if (!s)
    {
        report_error("Invalid block state");
        return QCOW2_ERROR_INVALID;
    }

    if (offset < 0 || offset >= bs->total_sectors * 512)
    {
        report_error("Offset out of range: %ld", offset);
        return QCOW2_ERROR_INVALID;
    }

    if (bytes <= 0)
    {
        report_error("Invalid read size: %ld", bytes);
        return QCOW2_ERROR_INVALID;
    }

    // 简单的读取实现 - 生产环境需要实现完整的簇查找
    FILE *f = fopen(bs->filename, "rb");
    if (!f)
    {
        report_error("Failed to open file for reading");
        return QCOW2_ERROR_IO;
    }

    if (fseek(f, offset, SEEK_SET) != 0)
    {
        fclose(f);
        return QCOW2_ERROR_IO;
    }

    int64_t result = fread(buf, 1, bytes, f);
    fclose(f);

    if (result != bytes)
    {
        report_error("Read incomplete: %ld/%ld bytes", result, bytes);
        return QCOW2_ERROR_IO;
    }

    return bytes;
}

// 写入数据
int qcow2_write(BlockDriverState *bs, int64_t offset, const uint8_t *buf, int64_t bytes)
{
    QCow2State *s = (QCow2State *)bs->opaque;
    if (!s)
    {
        report_error("Invalid block state");
        return QCOW2_ERROR_INVALID;
    }

    if (offset < 0 || offset >= bs->total_sectors * 512)
    {
        report_error("Offset out of range: %ld", offset);
        return QCOW2_ERROR_INVALID;
    }

    if (bytes <= 0)
    {
        report_error("Invalid write size: %ld", bytes);
        return QCOW2_ERROR_INVALID;
    }

    report_info("Writing %ld bytes at offset %ld", bytes, offset);

    // 简单的写入实现 - 生产环境需要实现完整的簇分配和写入
    FILE *f = fopen(bs->filename, "rb+");
    if (!f)
    {
        report_error("Failed to open file for writing");
        return QCOW2_ERROR_IO;
    }

    if (fseek(f, offset, SEEK_SET) != 0)
    {
        fclose(f);
        return QCOW2_ERROR_IO;
    }

    int64_t result = fwrite(buf, 1, bytes, f);
    fclose(f);

    if (result != bytes)
    {
        report_error("Write incomplete: %ld/%ld bytes", result, bytes);
        return QCOW2_ERROR_IO;
    }

    report_info("Write completed successfully");
    return bytes;
}

// 数据块查询 - 快速版本
std::list<DataBlockInfo *> qcow2_query_data_blocks_fast(BlockDriverState *bs)
{
    std::list<DataBlockInfo *> block_list;
    QCow2State *s = (QCow2State *)bs->opaque;

    if (!s)
    {
        report_error("Invalid block state");
        return block_list;
    }

    int64_t virtual_size = bs->total_sectors * 512;
    int64_t cluster_size = s->cluster_size;

    report_info("Fast querying data blocks: virtual_size=%ld, cluster_size=%d",
                virtual_size, cluster_size);

    // 直接分析 L1 和 L2 表来推断数据分布
    DataBlockInfo *current_block = nullptr;
    uint64_t current_offset = 0;

    // 遍历所有 L1 表项
    for (int l1_index = 0; l1_index < s->l1_size; l1_index++)
    {
        uint64_t l1_entry = s->l1_table[l1_index];
        if (!l1_entry)
        {
            // 整个 L2 表都未分配
            uint64_t segment_size = (uint64_t)s->l2_size << s->cluster_bits;
            if (current_block && current_block->cluster_type == QCOW2_CLUSTER_UNALLOCATED &&
                current_block->offset + current_block->length == current_offset)
            {
                current_block->length += segment_size;
            }
            else
            {
                current_block = (DataBlockInfo *)qemu_malloc(sizeof(DataBlockInfo));
                current_block->offset = current_offset;
                current_block->length = segment_size;
                current_block->file_offset = 0;
                current_block->cluster_type = QCOW2_CLUSTER_UNALLOCATED;
                current_block->allocated = false;
                block_list.push_back(current_block);
            }
            current_offset += segment_size;
            continue;
        }

        // L2 表已分配，需要读取 L2 表
        uint64_t l2_offset = l1_entry & L1E_OFFSET_MASK;
        uint64_t *l2_table = nullptr;

        FILE *f = fopen(bs->filename, "rb");
        if (f && qcow2_read_l2_table(s, f, l2_offset, &l2_table) == 0)
        {
            // 分析 L2 表中的所有簇
            for (int l2_index = 0; l2_index < s->l2_size; l2_index++)
            {
                uint64_t cluster_offset_entry = l2_table[l2_index];
                uint64_t offset = current_offset;
                int cluster_type;
                bool allocated;
                uint64_t file_offset = 0;

                if (cluster_offset_entry == 0) {
                    // 未分配簇
                    cluster_type = QCOW2_CLUSTER_UNALLOCATED;
                    allocated = false;
                } else if (cluster_offset_entry & QCOW_OFLAG_COMPRESSED) {
                    // 压缩簇
                    cluster_type = QCOW2_CLUSTER_COMPRESSED;
                    allocated = true;
                    file_offset = cluster_offset_entry & s->cluster_offset_mask;
                } else if (cluster_offset_entry & QCOW_OFLAG_COPIED) {
                    // 普通已分配簇
                    cluster_type = QCOW2_CLUSTER_NORMAL;
                    allocated = true;
                    file_offset = cluster_offset_entry & s->cluster_offset_mask;
                } else {
                    // 其他情况（可能是零簇或未分配）
                    // 在 QCow2 v3 中，这可能表示零簇
                    cluster_type = QCOW2_CLUSTER_UNALLOCATED;
                    allocated = false;
                }
                // 合并连续相同类型的块
                if (current_block &&
                    current_block->allocated == allocated &&
                    current_block->cluster_type == cluster_type &&
                    current_block->offset + current_block->length == offset)
                {

                    current_block->length += cluster_size;
                }
                else
                {
                    current_block = (DataBlockInfo *)qemu_malloc(sizeof(DataBlockInfo));
                    current_block->offset = offset;
                    current_block->length = cluster_size;
                    current_block->file_offset = file_offset;
                    current_block->cluster_type = cluster_type;
                    current_block->allocated = allocated;
                    block_list.push_back(current_block);
                }

                current_offset += cluster_size;
            }
            fclose(f);
        }
        else
        {
            // L2 表读取失败，假设整个段未分配
            uint64_t segment_size = (uint64_t)s->l2_size << s->cluster_bits;
            if (current_block && current_block->cluster_type == QCOW2_CLUSTER_UNALLOCATED &&
                current_block->offset + current_block->length == current_offset)
            {
                current_block->length += segment_size;
            }
            else
            {
                current_block = (DataBlockInfo *)qemu_malloc(sizeof(DataBlockInfo));
                current_block->offset = current_offset;
                current_block->length = segment_size;
                current_block->file_offset = 0;
                current_block->cluster_type = QCOW2_CLUSTER_UNALLOCATED;
                current_block->allocated = false;
                block_list.push_back(current_block);
            }
            current_offset += segment_size;
        }
    }

    report_info("Fast query found %lu data blocks", block_list.size());
    return block_list;
}

// 数据块查询 - 详细版本
std::list<DataBlockInfo *> qcow2_query_data_blocks(BlockDriverState *bs)
{
    return qcow2_query_data_blocks_fast(bs); // 简化实现
}

// 簇分配
uint64_t qcow2_alloc_clusters(BlockDriverState *bs, uint64_t size)
{
    QCow2State *s = (QCow2State *)bs->opaque;
    if (!s)
        return 0;

    FILE *f = fopen(bs->filename, "rb+");
    if (!f)
    {
        report_error("Failed to open file for cluster allocation");
        return 0;
    }

    // 获取文件大小
    if (fseek(f, 0, SEEK_END) != 0)
    {
        fclose(f);
        return 0;
    }

    uint64_t file_size = ftell(f);

    // 对齐到簇边界
    uint64_t cluster_offset = (file_size + s->cluster_size - 1) & ~(s->cluster_size - 1);

    // 扩展文件
    uint64_t new_size = cluster_offset + size;
    if (ftruncate(fileno(f), new_size) != 0)
    {
        report_error("Failed to truncate file to size %lu", new_size);
        fclose(f);
        return 0;
    }

    fclose(f);

    report_info("Allocated %lu clusters at offset 0x%lx", size / s->cluster_size, cluster_offset);
    return cluster_offset;
}

// 其他函数存根实现
int qcow2_allocate_clusters(BlockDriverState *bs, uint64_t offset, uint64_t size)
{
    return 0;
}

int qcow2_pwritev(BlockDriverState *bs, uint64_t offset, uint64_t bytes,
                  const uint8_t *buf, uint64_t flags)
{
    return qcow2_write(bs, offset, buf, bytes);
}

int qcow2_free_clusters(BlockDriverState *bs, uint64_t offset, uint64_t size)
{
    return 0;
}

int qcow2_update_cluster_refcount(BlockDriverState *bs, uint64_t cluster_offset, int add)
{
    return 0;
}

int qcow2_grow_l1_table(BlockDriverState *bs, int min_size)
{
    return 0;
}

int qcow2_grow_refcount_table(BlockDriverState *bs, int min_size)
{
    return 0;
}

int qcow2_alloc_cluster_at(BlockDriverState *bs, uint64_t offset, uint64_t *cluster_offset)
{
    return 0;
}

int qcow2_update_refcount(BlockDriverState *bs, uint64_t offset, uint64_t length, int add)
{
    return 0;
}

int qcow2_get_refcount(BlockDriverState *bs, uint64_t cluster_offset)
{
    return 0;
}

int qcow2_alloc_refcount_block(BlockDriverState *bs, uint64_t cluster_offset)
{
    return 0;
}

int qcow2_flush(BlockDriverState *bs)
{
    return 0;
}

int qcow2_check_consistency(BlockDriverState *bs)
{
    return 0;
}

uint64_t qcow2_virtual_size(BlockDriverState *bs)
{
    return bs->total_sectors * 512;
}

uint64_t qcow2_physical_size(BlockDriverState *bs)
{
    return get_file_size(bs->filename);
}

int qcow2_read_refcount_table(QCow2State *s, FILE *f)
{
    return 0;
}

int qcow2_write_l1_table(QCow2State *s)
{
    return 0;
}

int qcow2_write_refcount_table(QCow2State *s)
{
    return 0;
}

int qcow2_update_l2_entry(BlockDriverState *bs, uint64_t l2_offset, int l2_index,
                          uint64_t cluster_offset)
{
    return 0;
}